#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2021 wanghuacheng <wanghuacheng@wanghuacheng-PC>
#
# Distributed under terms of the MIT license.

"""

"""
import sys
import os
import argparse
import pandas as pd
import talib as ta
import numpy as np
import matplotlib.pyplot as plt
from talib.abstract import *
import datetime

date_range = [5, 13, 21]
len_date = len(date_range)

data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
merge_dir = os.path.join(data_dir, "history_k/merge/")
merge_dates = os.listdir(merge_dir)
md = max(merge_dates)
merge_latest_dir = os.path.join(merge_dir, md)
code_names = os.listdir(merge_latest_dir)
#print(codes)
#sys.exit(0)


today = datetime.datetime.today()
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('-c', '--code', type=str, default=None, help='code name')
parser.add_argument('-a', '--action', type=str, default="st1", help='code name')
parser.add_argument('-e', '--export', type=str, default=None, help='export')
parser.add_argument('--start_date', type=str, default='2020-01-01', help='export')
parser.add_argument('--end_date', type=str, default='2021-04-09', help='export')

ARGS = parser.parse_args()

today_str = today.strftime("%Y-%m-%d")

ALL_FEATURES = ['rsi_6', 'rsi_12', 'rsi_24', 'stoch_k', 'stoch_d', 'stoch_j', 'BBP', 'BLOW_HIGH', 'macd', 'macdsignal', 'macdhist', 'SMA_DIFF_5', 'SMA_DIFF_13', 'SMA_DIFF_21', 'SMA_DIFF_34', 'SMA_DIFF_55', 'SMA_DIFF_89', 'SMA_DIFF_144', 'SMA_DIFF_233', 'SMA_DIFF_377', 'SMA_5_13_DIFF', 'SMA_5_21_DIFF', 'SMA_5_34_DIFF', 'SMA_5_55_DIFF', 'SMA_5_89_DIFF', 'SMA_5_144_DIFF', 'SMA_5_233_DIFF', 'SMA_5_377_DIFF', 'SMA_13_21_DIFF', 'SMA_13_34_DIFF', 'SMA_13_55_DIFF', 'SMA_13_89_DIFF', 'SMA_13_144_DIFF', 'SMA_13_233_DIFF', 'SMA_13_377_DIFF', 'SMA_21_34_DIFF', 'SMA_21_55_DIFF', 'SMA_21_89_DIFF', 'SMA_21_144_DIFF', 'SMA_21_233_DIFF', 'SMA_21_377_DIFF', 'SMA_34_55_DIFF', 'SMA_34_89_DIFF', 'SMA_34_144_DIFF', 'SMA_34_233_DIFF', 'SMA_34_377_DIFF', 'SMA_55_89_DIFF', 'SMA_55_144_DIFF', 'SMA_55_233_DIFF', 'SMA_55_377_DIFF', 'SMA_89_144_DIFF', 'SMA_89_233_DIFF', 'SMA_89_377_DIFF', 'SMA_144_233_DIFF', 'SMA_144_377_DIFF', 'SMA_233_377_DIFF', 'peTTM_377_SMA_diff', 'pbMRQ_377_SMA_diff', 'psTTM_377_SMA_diff', 'pcfNcfTTM_377_SMA_diff', 'VOL_SMA_DIFF_5', 'VOL_SMA_DIFF_13', 'VOL_SMA_DIFF_21', 'VOL_SMA_DIFF_34', 'VOL_SMA_DIFF_55', 'VOL_SMA_DIFF_89', 'VOL_SMA_DIFF_144', 'VOL_SMA_DIFF_233', 'VOL_SMA_DIFF_377']

def gold_cross(last_diff, today_diff):
    return last_diff < 0 and today_diff >= 0

def dead_cross(last_diff, today_diff):
    return last_diff > 0 and today_diff <= 0

def is_top(df, index):
    close = df.at[index, 'close']
    high = df.at[index, 'high']
    dflag = True
    for i in range(1, 7):
        # print(df.loc[index, 'close'])
        if close >= df.at[index + i, 'close'] and close >= df.at[index - i, 'close']:
            continue
        else:
            dflag = False
            break

    cp6 = df.at[index - 6, 'close']
    ca6 = df.at[index + 6, 'close']


    hp1 = df.at[index - 1, 'high']
    ha1 = df.at[index + 1, 'high']
    tfx = high > hp1 and high > ha1

    climit = close * 0.97

    return tfx and dflag and cp6 < climit and ca6 < climit



    return dflag

def is_bottom(df, index):
    close = df.at[index, 'close']
    low = df.at[index, 'low']

    dflag = True
    for i in range(1, 7):
        # print(df.loc[index, 'close'])
        if close <= df.at[index + i, 'close'] and close <= df.at[index - i, 'close']:
            continue
        else:
            dflag = False
            break
    cp6 = df.at[index - 6, 'close']
    ca6 = df.at[index + 6, 'close']
    la1 = df.at[index + 1, 'low']
    lp1 = df.at[index - 1, 'low']
    lfx = low < la1 and low < lp1

    climit = close * 1.03
    return lfx and dflag and cp6 > climit and ca6 > climit


def find_st1(df):
    N = df.shape[0]
    if N < 500:
        return None

    features = []

    # rsi
    for d in [6, 12, 24]:
        dfn = 'rsi_' + str(d)
        features.append(dfn)
        df[dfn] = RSI(df, timeperiod = d) / 100


    matype = 0
    # slowk      slowd
    stoch = STOCH(df, fastk_period=9, slowk_matype=matype, slowk_period=3, slowd_period=3, slowd_matype=matype)
    df['stoch_k'] = stoch['slowk'] / 100
    df['stoch_d'] = stoch['slowd'] / 100
    df['stoch_j'] = (stoch['slowk'] * 3 - stoch['slowd'] * 2)/100
    kdj_features = ['stoch_k', 'stoch_d', 'stoch_j']
    features += kdj_features



    features += ['BBP', "BLOW_HIGH"]
    bands = BBANDS(df, timeperiod = 20, nbdevup = 2.0, nbdevdn = 2.0, matype = 0)

    df['BBP'] = (df.close - bands.lowerband) / (bands.upperband - bands.lowerband)
    df['BLOW_HIGH'] = bands.upperband / bands.lowerband



    # print(stoch)
    macd_df = MACD(df, fastperiod=12, slowperiod=26, signalperiod=9)
    # macd  macdsignal  macdhist
    # print(macd_df)
    df = pd.concat([df, macd_df, stoch], axis = 1)
    macd_features = ['macd','macdsignal','macdhist']
    features += macd_features
    # close features
    sma_close_diff_names = []
    for d in date_range:
        diff_name = 'SMA_DIFF_' + str(d)
        sma_name = 'SMA_' + str(d)
        sma_close_diff_names.append(diff_name)

        df[sma_name] = SMA(df, timeperiod = d)
        df[diff_name] = (df['close'] - df[sma_name]) / df[sma_name]


    sma_diff_names = []
    for i in range(0, len_date - 1):
        for j in range(i + 1, len_date):
            diff_name = "SMA_" + str(date_range[i]) + "_" + str(date_range[j]) + "_DIFF"
            sma_diff_names.append(diff_name)
            df[diff_name] = (df['SMA_' + str(date_range[i])] - df['SMA_' + str(date_range[j])]) / df['SMA_' + str(date_range[j])]

    sma_fes = ['SMA_' + str(d) for d in date_range]

    features += sma_close_diff_names + sma_diff_names
    # print(features)


    ## val features .............
    #val_features = ["peTTM", "pbMRQ", "psTTM", "pcfNcfTTM"]
    #val_diff_features = []
    #for vf in val_features:
    #    df[vf] = df[vf].astype('float')
    #    sma_name = vf + '_377_SMA'
    #    df[sma_name] = SMA(df, timeperiod = 377, price = vf)
    #    # vf_mean = df[vf].mean()
    #    # print(vf_mean)
    #    vf_diff_name = sma_name + "_diff"
    #    val_diff_features.append(vf_diff_name)
    #    df[vf_diff_name] = (df[vf] - df[sma_name])/ df[sma_name]


    # volume
    sma_vol_diff_names = []
    for d in date_range:
        diff_name = 'VOL_SMA_DIFF_' + str(d)
        sma_name = 'VOL_SMA_' + str(d)
        sma_vol_diff_names.append(diff_name)

        df[sma_name] = SMA(df, timeperiod = d, price = 'volume')
        df[diff_name] = (df['volume'] - df[sma_name]) / df[sma_name]




    #features = features + val_diff_features + sma_vol_diff_names
    features = features + sma_vol_diff_names
    # print(features)
    df = df.fillna(method="bfill").fillna(method="ffill").fillna(0)
    df.replace([np.inf, -np.inf], 0, inplace = True)
    df['dt'] = pd.to_datetime(df.date)
    if ARGS.export:
        #fvals = (df.loc[N-9 : N - 1, ['date', 'code'] + features].values)
        fvals = (df[(df.dt >= ARGS.start_date) & (df.dt <= ARGS.end_date)][['date', 'code'] + features].values)
        return fvals
    else:
        info = []
        for index, row in df.iterrows():
            if index < 500:
                continue
            if index > N - 30:
                break

            dflag = is_bottom(df, index)
            tflag = is_top(df, index)
            if dflag:
                target = 1
            if tflag:
                target = 0

            if dflag or tflag:
                fvals = (df.loc[index-1, ['date', 'code'] + features].values)
                # print(dict(zip(features, fvals)))
                # print(row.date, fvals, target)
                # yield fvals, target
                info.append((fvals, target))
        # sys.exit(0)

        return info







def get_date_diff(d1, d2):
    dd1 = datetime.datetime.strptime(d1, '%Y-%m-%d')
    dd2 = datetime.datetime.strptime(d2, '%Y-%m-%d')
    return (dd1 - dd2).days

if __name__ == '__main__':
    code = ARGS.code
    action = ARGS.action
    fun = eval("find_"+ action)
    outfile = action+"_features.csv"
    codes = None
    if code == 'hs300':
        code_df = pd.read_csv(os.path.join(data_dir , 'hs300.csv'))
        codes = set(code_df.code.tolist())
        outfile = code + "_features.csv"
    if ARGS.export:
        outfile = "export_features.csv"
    # print(codes)
    # print(code_names)

    with open(outfile, "w") as fw:
        if ARGS.export:
            fw.write("date,code," + ",".join(ALL_FEATURES) + "\n")
        else:
            fw.write(",".join(ALL_FEATURES) + ",target\n")
        for cn in code_names:
            code = cn[:9]
            # print(code)
            if codes:
                if code in codes:
                    print(code)
                else:
                    continue
            code_name = cn[10:-4]
            cnfile = os.path.join(merge_latest_dir, cn)
            #print(cnfile)
            df = pd.read_csv(cnfile)
            info = fun(df)
            if not info is None:
                if not ARGS.export:
                    for fvals, target in info:
                        fw.write(",".join(map(str, fvals))+","+str(target) + "\n")
                    fw.flush()
                else:
                    for line in info:
                        fw.write(",".join(map(str, line)) + "\n")
                        fw.flush()
